package com.bungarus.ratelimit;

import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.core.Authentication;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import javax.servlet.http.HttpServletRequest;
import java.security.Principal;

@Component
@Aspect
public class RateLimitAspect {
    private LimiterTokenParser tokenParser;
    private RateLimitManager rateLimitManager;

    @Autowired
    public RateLimitAspect(LimiterTokenParser tokenParser, RateLimitManager rateLimitManager) {
        this.tokenParser = tokenParser;
        this.rateLimitManager = rateLimitManager;
    }

    @Pointcut("@target(org.springframework.web.bind.annotation.RestController)")
    public void annotationController() {}

    @Pointcut("execution(* *..*(javax.servlet.http.HttpServletRequest,..))")
    public void methodWithRequest() {}

    @Pointcut("execution(* *..*(org.springframework.security.core.Authentication,..))")
    public void methodWithPrincipal() {}

    @Pointcut("@annotation(annotationLimiter)")
    public void annotationLimiterDef(Limiter annotationLimiter) {
    }

    @Before("annotationController() && annotationLimiterDef(annotationLimiter) && methodWithRequest() && args(request,..)")
    public void beforeAdviceWithRequest(JoinPoint joinPoint, Limiter annotationLimiter, HttpServletRequest request) throws RateLimitException {
        String id;
        //Precondition cannot be matched
        if(annotationLimiter.request().length() == 0 && annotationLimiter.header().length() == 0) {
            throw new RateLimitException("Principal in request is needed.");
        }
        else if(annotationLimiter.request().length() > 0) {
            LimiterIdentity identity = (LimiterIdentity)request.getAttribute(annotationLimiter.request());
            if(null == identity) {
                throw new RateLimitException("Principal in HttpServletRequest is needed.");
            }
            else {
                id = identity.identity();
            }
        }
        else {
            String authHeader = request.getHeader(annotationLimiter.header());
            if(StringUtils.isEmpty(authHeader)) {
                throw new RateLimitException("Authentication token in header is required.");
            }
            id = tokenParser.parse(authHeader);
        }
        if(id.length() == 0) {
            throw new RateLimitException("Principal is invalid.");
        }
        if(-1 == rateLimitManager.pass(id + ":" + joinPoint.getSignature().getName(), annotationLimiter.limit(), annotationLimiter.unit())) {
            throw new RateLimitExceedException("Exceed the request maximum threshold.", annotationLimiter.limit(), annotationLimiter.unit());
        }
    }

    //Using spring security to get the identity
    @Before("annotationController() && annotationLimiterDef(annotationLimiter) && methodWithPrincipal() && args(authentication,..)")
    public void beforeAdviceWithSecurity(JoinPoint joinPoint, Limiter annotationLimiter, Authentication authentication) throws RateLimitException {
        if(null == authentication || !authentication.isAuthenticated()) {
            throw new RateLimitException("Authenticated principal is needed.");
        }
        if(-1 == rateLimitManager.pass(authentication.getPrincipal() + ":" + joinPoint.getSignature().getName(), annotationLimiter.limit(), annotationLimiter.unit())) {
            throw new RateLimitExceedException("Exceed the request maximum threshold.", annotationLimiter.limit(), annotationLimiter.unit());
        }
    }
}
